"""
PSO test using the Synchronous Handler
"""
#from playdoh03 import *
from ..rpc import *
from ..debugtools import *
from ..gputools import *
from ..syncjobhandler import *
from ..pool import *
import numpy as np, numpy.random as nprnd, math, time
import threading

__all__ = ['Optimization']

class Optimization(SyncJobHandler):
    def __init__(self, fitness = None, **params):
        SyncJobHandler.__init__(self)
        self.set_fitness(fitness, **params)
        self.set_no_machines()
        self._status = None
    
    @staticmethod
    def default_values(**opt_info):
        """
        To be overloaded by optimization algorithms.
        """
        return opt_info
    
    def set_fitness(self, fitness, **params):
        self.fitness_fun = fitness
        self.params = params
        self.optparams = OptParams(**params)
    
    def set_opt_info(self, **opt_info):
        """
        optinfo is a dictionary with the following keys:
            _iterations
            _cpu
            _gpu
            ...
        These keys must be set (with default values at least) in the Interface (minimize and maximimize functions)
        This function returns optresults, a dictionary with the best fitness, the best params,
        and other information about the optimization (histogram of the fitness for instance, etc.,
        to be defined by the optimization algorithm).
        """
        self.opt_info = opt_info
        self.iterations = opt_info['_iterations']
        self.particles = opt_info['_particles']
        self.type = opt_info['_type'] # CPU or GPU
        self.units = opt_info['_units'] # number of CPUs/GPUs
        
        self.opt_info=self.default_values(**opt_info) # this should work even if default_values if static
    
    def get_class(self):
        return self.__class__#.__name__
    
    def set_no_machines(self):
        self.machines = None
        self.machine_index = 0
        self._other_machines = None
    
    def set_machines(self, machines):
        """
        Sets the list of the machines used to run the optimization. The current machine name
        is self.get_machine_id() and it must be an element of self.machines.
        """
        if machines is None or type(machines) is not list or len(machines) == 0:
            self.set_no_machines()
            return
        
        machines.sort()
        self.machines = machines
        
        # Sets other machines
        if (self.this_machine, self.port) not in self.machines:
            raise Exception("The current machine '%s' must be in 'self.machines'" % self.this_machine)
        self.machine_index = self.machines.index((self.this_machine, self.port))
        self._other_machines = self.machines[:] # creates a copy of the list
        self._other_machines.remove((self.this_machine, self.port))
        
        log.debug("this machine : %s" % self.this_machine)
        log.debug("this machine index : %s" % self.this_machine_index)
        log.debug("all machines : %s" % self.all_machines)
        log.debug("other machines : %s" % self.other_machines)
        
    def _get_this_machine(self):
        return self.machine_id
    this_machine = property(_get_this_machine)
    
    def _get_this_machine_index(self):
        return self.machine_index
    this_machine_index = property(_get_this_machine_index)
    
    def _get_all_machines(self):
        """
        self.machines read_only
        """
        return self.machines
    all_machines = property(_get_all_machines)
    
    def _get_other_machines(self):
        return self._other_machines
    other_machines = property(_get_other_machines)

    def _get_sub_fitness(self, fitness_fun, sub_params):
        return fitness_fun(**sub_params)

    def get_fitness(self, param_matrix = None):
        """
        Computes the fitness from the parameters matrix.
        self.fitness is updated here.
        Uses several CPUs or GPUs.
        """
        if param_matrix is None:
            param_matrix = self.X
        
        log.debug("Computing fitness")
        if self.type == 'CPU' and self.units == 1:
            param_values = self.optparams.get_param_values(param_matrix)
            results = self.fitness_fun(**param_values)
        else:
            sub_matrices = self.optparams.split_matrix(param_matrix, self.units)
            sub_params = self.optparams.get_split_param_values(sub_matrices)
            
#            p = CustomPool(self.units, self.type)
            self._pool.reset()
            self._pool.map(self.fitness_fun, **sub_params) 
#            p.map(self._get_sub_fitness, [fitness_fun]*self.units, sub_params) # TODO
            results = self._pool.get_results() # the list of the fitness for each group
#            p.close()
#            del p
            
        fit = np.hstack(results)
        self.fitness = fit
        return fit

    def contact(self, machines, method, *argss, **kwdss):
        """
        Contacts machines by calling RPC method on them with given args and kwds.
        Can be used with self.all_machines or self.other_machines. 
        """
        
        log.debug("Contacting machines '%s' with method '%s'" % (str(machines), method))
        
        results = self.remote_call(machines, method, *argss, **kwdss)
            
        for result in results:
            if isinstance(result, Exception):
                raise result
            
        return results

    """
    Methods to be implemented by deriving classes
    """
    def initialize(self):
        """
        Initializes the optimization algorithm
        """
        log.warn("'self.initialize' method is not implemented")
        pass

    def iterate(self):
        """
        Performs the iteration step of the algorithm.
        The following variables are available :
            * self.param_values = initial parameters # current param values
            * self.X = corresponding initial matrix # current matrix
            * self.fitness = fitness values of initial parameters # current fitness
            * self.iteration = 0 # current iteration
            * self.opt_info = dict(algorithm-specific optimization parameters) # static
        The variable self.results must be filled by this method.
        The following methods can be used :
            * self.contact : contacts other machines 
            * self.get_fitness(matrix) computes the fitness of the given matrix
            * self.this_machine, self.this_machine_index, other_machines, all_machines in self.contact()
        """
        log.warn("'self.iterate' method is not implemented")
        pass

    def _optimize(self):
        """
        Optimizes the fitness function with respect to given params.
        """
        
        # Initializes pool of processes
        self._pool = CustomPool(self.units, self.type, close_pool = True)
        
        self.param_values = self.optparams.get_initial_param_values(self.particles)
        self.X = self.optparams.get_param_matrix(self.param_values)
        self.fitness = self.get_fitness()
        self.results = dict()
        
        # Optimization initialization
        log.info("Initializing optimization")
        self.initialize()
        self._status = "initialized"
        
        # Iterations
        for self.iteration in xrange(self.iterations):
            log.info("Iteration %d/%d" % (self.iteration+1, self.iterations))
            self.iterate()
            self._status = "iteration %d/%d" % (self.iteration+1, self.iterations)
        
        self._status = "finished"
#        self._pool.close()
        
        # results is a dictionary with at least the following keys :
        # fitness, 'param_key' for each key, with the best parameters for each group
        return self.results
    
    def optimize(self):
        self.thread = threading.Thread(target=self._optimize)
        self.thread.start()
        
    def get_status(self):
        return self._status
        
    def get_results(self):
        self.thread.join()
        return self.results
    
    @staticmethod
    def finalize(results_machines):
        """
        Called at the end of the optimization by the client.
        Retrieves the results of all the machines and return
        a global result (best_fit, best_position, best_fits)
        results_machines is a list of results
        """
        log.warn("'self.finalize' method is not implemented")
        pass
    



class OptParams(object):
    """Internal class used to manipulate model fitting parameters.
    It basically handles conversion between parameter dictionaries and arrays.
    
    Initialized with arguments:
    
    ``**params``
        Parameters list ``param_name=[bound_min, min_ max, bound_max]``
        
    **Methods**
    
    .. method:: get_initial_param_values(N)
    
        Returns initial parameter values sampled uniformly within the parameter 
        interval given in the constructor of the class. ``N`` is the number of neurons.
        The result is a dictionary ``{param_name=values}`` where values is a vector of values.
    
    .. method:: set_constraints()
    
        Returns the constraints for each parameter. The result is (min_values, max_values)
        where each variable is a vector containing the minimum and maximum values for each parameter.
    
    .. method:: get_param_values(X)
    
        Converts an array containing parameter values into a dictionary.
    
    .. method:: get_param_matrix(param_values)
    
        Converts a dictionary containing parameter values into an array.
    """
    def __init__(self, **params):
        self.params = params
        self.param_names = params.keys()
        self.param_names.sort()
        self.param_count = len(params)
        
    def get_initial_param_values(self, N):
        """
        Samples random initial param values around default values
        """
        random_param_values = {}
        for key, value in self.params.iteritems():
            if np.isscalar(value):
                value = [value]
#            if len(value) == 1:
#                # One default value, no boundary counditions on parameters
#                random_param_values[key] = value[0]*(1+.5*randn(N))
            elif len(value) == 2:
                # One default interval, no boundary counditions on parameters
                random_param_values[key] = value[0] + (value[1]-value[0])*nprnd.rand(N)
#            elif len(value) == 3:
#                # One default value, value = [min, init, max]
#                random_param_values[key] = value[1]*(1+.5*randn(N))
            elif len(value) == 4:
                # One default interval, value = [min, init_min, init_max, max]
                random_param_values[key] = value[1] + (value[2]-value[1])*nprnd.rand(N)
        return random_param_values

    def set_constraints(self):
        """
        Returns constraints of a given model
        returns min_values, max_values
        min_values is an array of length p where p is the number of parameters
        min_values[i] is the minimum value for parameter i
        """
        min_values = []
        max_values = []
        param_names = self.param_names
        p = self.param_count
        for key in param_names:
            value = self.params[key]
            # No boundary conditions if only two values are given
            if len(value) == 2:
                min_values.append(-inf)
                max_values.append(inf)
            else:
                min_values.append(value[0])
                max_values.append(value[-1])
        
        min_values = array(min_values)
        max_values = array(max_values)
        
        return min_values, max_values

    def get_param_values(self, X):
        """
        Converts a matrix containing param values into a dictionary
        """
        param_values = {}
        if X.ndim == 1:
            X = X.reshape((-1,1))
        for i in range(len(self.param_names)):
            param_values[self.param_names[i]] = X[i,:]
        return param_values
    
    def get_param_matrix(self, param_values):
        """
        Converts a dictionary containing param values into a matrix
        """
        p = self.param_count
        # Number of parameter values (number of particles)
        n = len(param_values[self.param_names[0]])
        X = np.zeros((p, n))
        for i in range(p):
            X[i,:] = param_values[self.param_names[i]]
        return X
    
    def split_matrix(self, X, k, axis=1):
        """
        Splits a matrix into k matrices (along the given axis, default j-axis)
        """
        N = X.shape[axis]
        size = [N/k for _ in xrange(k)]
        size[-1] = N-sum(size[:-1])
        assert sum(size) == N
        return np.split(X, np.cumsum(size)[:-1], axis=axis)
        
    def get_split_param_values(self, sub_matrices):
        """
        Returns split param values from sub matrices obtained with split_matrix
        result[key] is the list of the param values for the given parameter and for every subgroup
        """
        split_params = dict()
        sub_params = [self.get_param_values(sub_matrix) for sub_matrix in sub_matrices]
        for key in self.param_names:
            split_params[key] = [sub_params[i][key] for i in xrange(len(sub_matrices))]
        return split_params

